import fitz
import numpy as np
from PIL import Image
import io
import os
import easyocr
import glob
import re
import logging
from collections import Counter
from datetime import datetime


# 配置日志系统
def setup_logging():
    """设置日志配置"""
    # 创建logs目录
    if not os.path.exists('logs'):
        os.makedirs('logs')

    # 生成带时间戳的日志文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_filename = f'logs/pdf_processing_{timestamp}.log'

    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_filename, encoding='utf-8'),
            logging.StreamHandler()  # 同时输出到控制台
        ]
    )

    return logging.getLogger(__name__)


# 初始化日志
logger = setup_logging()


def ocr_pdf_easyocr(pdf_path):
    """使用EasyOCR识别PDF文字（修复版）"""
    try:
        if not os.path.exists(pdf_path):
            logger.error(f"文件不存在: {pdf_path}")
            return f"错误: 文件 {pdf_path} 不存在"

        logger.info("初始化EasyOCR...")
        # 初始化OCR阅读器
        reader = easyocr.Reader(['ch_sim', 'en'])  # 中文和英文

        doc = fitz.open(pdf_path)
        full_text = ""
        total_pages = len(doc)

        logger.info(f"开始处理PDF: {pdf_path}")
        logger.info(f"总页数: {total_pages}")
        logger.info("-" * 50)

        for page_num in range(total_pages):
            logger.info(f"正在处理第 {page_num + 1} 页...")

            # 获取PDF页面
            page = doc.load_page(page_num)

            # 设置分辨率
            mat = fitz.Matrix(2.0, 2.0)  # 使用缩放因子而不是DPI
            pix = page.get_pixmap(matrix=mat)

            # 方法1：使用PIL作为中介（推荐）
            try:
                # 将pixmap转换为PIL Image
                img_data = pix.tobytes("ppm")
                pil_image = Image.open(io.BytesIO(img_data))

                # 转换为RGB模式（确保3通道）
                if pil_image.mode != 'RGB':
                    pil_image = pil_image.convert('RGB')

                # 将PIL Image转换为numpy数组
                img_array = np.array(pil_image)

            except Exception as img_error:
                logger.error(f"图片转换错误: {img_error}")
                # 方法2：直接处理（备选）
                try:
                    # 获取图片尺寸和数据
                    width, height = pix.width, pix.height
                    # 直接使用pixmap的samples数据
                    img_array = np.frombuffer(pix.samples, dtype=np.uint8)
                    img_array = img_array.reshape(height, width, 3)
                except Exception as direct_error:
                    logger.error(f"直接转换也失败: {direct_error}")
                    continue

            # 检查数组形状
            logger.info(f"图片尺寸: {img_array.shape}")

            # 识别文字
            try:
                results = reader.readtext(img_array)

                page_text = ""
                text_count = 0

                for (bbox, text, confidence) in results:
                    if confidence > 0.3:  # 置信度阈值
                        page_text += text + " "
                        text_count += 1

                full_text += f"\n--- 第 {page_num + 1} 页 ---\n{page_text}\n"
                logger.info(f"✓ 第 {page_num + 1} 页识别完成，找到 {text_count} 个文本")

            except Exception as ocr_error:
                logger.error(f"OCR识别错误: {ocr_error}")
                full_text += f"\n--- 第 {page_num + 1} 页 ---\n[识别错误: {ocr_error}]\n"

        doc.close()
        logger.info("-" * 50)
        logger.info("PDF处理完成！")

        return full_text if full_text.strip() else "未识别到任何文字"

    except Exception as e:
        logger.error(f"OCR处理异常: {str(e)}")
        return f"错误: {str(e)}"


# 更稳定的版本（使用PIL处理）
def ocr_pdf_easyocr_stable(pdf_path, dpi=150):
    """更稳定的EasyOCR PDF识别版本"""
    try:
        if not os.path.exists(pdf_path):
            logger.error(f"文件不存在: {pdf_path}")
            return f"错误: 文件 {pdf_path} 不存在"

        logger.info("初始化EasyOCR...")
        reader = easyocr.Reader(['ch_sim', 'en'])

        doc = fitz.open(pdf_path)
        full_text = ""
        total_pages = len(doc)

        logger.info(f"处理PDF: {pdf_path} (共{total_pages}页)")
        logger.info("-" * 50)

        for page_num in range(total_pages):
            logger.info(f"处理第 {page_num + 1} 页...")

            try:
                # 使用更安全的方式转换图片
                page = doc.load_page(page_num)

                # 方法：先保存为临时图片，再读取
                temp_img_path = f"temp_page_{page_num}.png"

                # 渲染页面为图片
                pix = page.get_pixmap(matrix=fitz.Matrix(dpi / 72, dpi / 72))
                pix.save(temp_img_path)

                # 使用EasyOCR识别图片文件
                results = reader.readtext(temp_img_path)

                # 提取文字
                page_text = ""
                for (bbox, text, confidence) in results:
                    if confidence > 0.2:  # 降低置信度阈值
                        page_text += text + " "

                full_text += f"\n--- 第 {page_num + 1} 页 ---\n{page_text}\n"
                logger.info(f"✓ 识别到 {len(results)} 个文本区域")

                # 删除临时文件
                os.remove(temp_img_path)

            except Exception as page_error:
                logger.error(f"第 {page_num + 1} 页处理失败: {page_error}")
                full_text += f"\n--- 第 {page_num + 1} 页 ---\n[处理失败]\n"

                # 清理临时文件
                if os.path.exists(f"temp_page_{page_num}.png"):
                    os.remove(f"temp_page_{page_num}.png")

        doc.close()
        return full_text

    except Exception as e:
        logger.error(f"稳定版OCR处理异常: {str(e)}")
        return f"错误: {str(e)}"


# 最简单的版本（最高成功率）
def ocr_pdf_simple(pdf_path):
    """最简单的EasyOCR PDF识别版本"""
    try:
        if not os.path.exists(pdf_path):
            logger.error(f"文件不存在: {pdf_path}")
            return f"错误: 文件 {pdf_path} 不存在"

        logger.info("初始化EasyOCR...")
        reader = easyocr.Reader(['ch_sim', 'en'])

        doc = fitz.open(pdf_path)
        full_text = ""

        for page_num in range(len(doc)):
            logger.info(f"处理第 {page_num + 1} 页...")

            # 将每一页保存为临时图片
            page = doc.load_page(page_num)
            pix = page.get_pixmap(matrix=fitz.Matrix(150 / 72, 150 / 72))  # 150 DPI
            temp_img = f"temp_{page_num}.png"
            pix.save(temp_img)

            # 直接识别图片文件
            results = reader.readtext(temp_img)

            page_text = ' '.join([text for (bbox, text, confidence) in results if confidence > 0.2])
            full_text += f"\n--- 第 {page_num + 1} 页 ---\n{page_text}\n"

            # 删除临时文件
            os.remove(temp_img)

        doc.close()
        return full_text

    except Exception as e:
        logger.error(f"简单版OCR处理异常: {str(e)}")
        return f"错误: {str(e)}"


def find_most_frequent_robust(text1, text2, text3):
    """健壮版本：处理各种特殊情况"""

    # 清理文本（移除多余空格和特殊字符）
    def clean_text(text):
        if not isinstance(text, str):
            return ""
        # 替换全角字符为半角
        text = text.replace('，', ',').replace('。', '.').replace('？', '?')
        # 移除多余空格
        text = ' '.join(text.split())
        return text

    text1_clean = clean_text(text1)
    text2_clean = clean_text(text2)
    text3_clean = clean_text(text3)

    pattern = r'(?<!\S)\d{18}(?!\S)'  # 匹配前后为空格或边界的18位数字

    # 提取18位数字
    numbers1_18 = re.findall(pattern, text1_clean)
    numbers2_18 = re.findall(pattern, text2_clean)
    numbers3_18 = re.findall(pattern, text3_clean)

    # 提取7位数字
    pattern_7 = r'(?<!\S)\d{7}(?!\S)'  # 匹配前后为空格或边界的7位数字
    numbers1_7 = re.findall(pattern_7, text1_clean)
    numbers2_7 = re.findall(pattern_7, text2_clean)
    numbers3_7 = re.findall(pattern_7, text3_clean)

    all_numbers = numbers1_18 + numbers2_18 + numbers3_18 + numbers1_7 + numbers2_7 + numbers3_7

    logger.info(f"清理后文本1 -> 18位数字: {numbers1_18}, 7位数字: {numbers1_7}")
    logger.info(f"清理后文本2 -> 18位数字: {numbers2_18}, 7位数字: {numbers2_7}")
    logger.info(f"清理后文本3 -> 18位数字: {numbers3_18}, 7位数字: {numbers3_7}")
    logger.info(f"所有匹配到的数字: {all_numbers}")

    if not all_numbers:
        logger.warning("未找到任何符合条件的数字")
        return None

    counter = Counter(all_numbers)
    logger.info(f"数字统计详情: {counter.most_common()}")

    most_common = counter.most_common(1)[0][0]
    logger.info(f"最频繁出现的数字: {most_common}")

    return most_common


if __name__ == "__main__":
    logger.info("=== PDF处理程序开始运行 ===")

    pdf_paths = glob.glob('*.pdf')
    logger.info(f"找到 {len(pdf_paths)} 个PDF文件: {pdf_paths}")

    if not pdf_paths:
        logger.warning("当前目录下未找到PDF文件")

    for pdf_path in pdf_paths:
        logger.info(f"\n{'=' * 60}")
        logger.info(f"开始处理文件: {pdf_path}")
        logger.info(f"{'=' * 60}")

        try:
            # 识别PDF文字
            logger.info("使用第一种OCR方法...")
            text1 = ocr_pdf_easyocr(pdf_path)

            logger.info("使用第二种OCR方法...")
            text2 = ocr_pdf_easyocr_stable(pdf_path)

            logger.info("使用第三种OCR方法...")
            text3 = ocr_pdf_simple(pdf_path)

            logger.info("分析识别结果...")
            number = find_most_frequent_robust(text1, text2, text3)

            logger.info(f"最终识别结果: {number}")

            if number:
                # 获取文件目录和扩展名
                dir_name = os.path.dirname(pdf_path)
                file_ext = os.path.splitext(pdf_path)[1]

                # 构建新文件名
                new_filename = f"{number}{file_ext}"
                new_path = os.path.join(dir_name, new_filename) if dir_name else new_filename

                # 重命名文件
                try:
                    os.rename(pdf_path, new_path)
                    logger.info(f"重命名成功: {pdf_path} -> {new_path}")
                except Exception as e:
                    logger.error(f"重命名失败: {e}")
            else:
                logger.warning("未识别到符合条件的数字，跳过重命名")

        except Exception as e:
            logger.error(f"处理文件 {pdf_path} 时发生异常: {e}")
            continue

    logger.info("=== PDF处理程序运行结束 ===")